import torch
import torch.nn.functional as F
import jieba
from models.rnn_model import TextRNN

import pickle

# -----------------------------
# 加载词表和模型参数
# -----------------------------

with open('dataset/vocab.pkl', 'rb') as f:
    vocab = pickle.load(f)

vocab_size = len(vocab)
embed_dim = 100
hidden_size = 128
num_classes = 2

model = TextRNN(vocab_size, embed_dim, hidden_size, num_classes)
model.load_state_dict(torch.load("model.pth", map_location=torch.device('cpu')))
model.eval()

# -----------------------------
# 句子预处理函数
# -----------------------------
def preprocess(text, vocab, max_len=20):
    tokens = list(jieba.cut(text))
    indices = [vocab.get(token, vocab['<unk>']) for token in tokens]
    if len(indices) < max_len:
        indices += [vocab['<pad>']] * (max_len - len(indices))
    else:
        indices = indices[:max_len]
    return torch.tensor(indices).unsqueeze(0)  # shape: [1, max_len]

# -----------------------------
# 单句测试
# -----------------------------
text = "剧情拖沓，演员表现一般。"
input_tensor = preprocess(text, vocab)  # shape: [1, max_len]
with torch.no_grad():
    output = model(input_tensor)  # logits
    probs = F.softmax(output, dim=1)
    pred = torch.argmax(probs, dim=1).item()

print(f"【输入句子】：{text}")
print(f"【预测类别】：{pred}（{'正面' if pred == 1 else '负面'}）")
print(f"【类别概率】：{probs.squeeze().tolist()}")
